#!/usr/bin/python3
#-*- coding: utf-8 -*-

import asyncio
import hexdump

class tcpproxy:
    def __init__(self, remote_ip, remote_port, target_ip=None, target_port=None):
        self.__remote_ip = remote_ip
        self.__remote_port = remote_port

        self.__target_ip = target_ip
        self.__target_port = target_port
        
    async def __copy_data_thread(self, info, reader, writer):
        addr = writer.get_extra_info('peername')
    
        while True:
            chunk = await reader.read(1024)
            if not chunk:
                break
            print("{0}:{1}".format(info, addr))
            hexdump.hexdump(chunk)
            writer.write(chunk)
            await writer.drain()

        print("close sock[{0}]:{1}".format(info, addr))
        writer.close()

    async def accept_handle(self, a_reader, a_writer):
        c_reader, c_writer = await asyncio.open_connection(self.__remote_ip, self.__remote_port)

        if self.__target_ip != None:
            await self.__http_connect_forward(c_reader, c_writer)

        dltasks = set()
        dltasks.add(asyncio.ensure_future(self.__copy_data_thread('c->a', c_reader, a_writer)))
        dltasks.add(asyncio.ensure_future(self.__copy_data_thread('a->c', a_reader, c_writer)))
        dones, dltasks = await asyncio.wait(dltasks, return_when=asyncio.FIRST_COMPLETED)

    async def __http_connect_forward(self, c_reader, c_writer):
        print("__http_connect_forward")
        cmd_connect = "CONNECT {0}:{1} HTTP/1.1\r\n\r\n".format(self.__target_ip, self.__target_port)
        print("--> {0}".format(str(cmd_connect)))
        c_writer.write(str.encode(cmd_connect))
        await c_writer.drain()
        
        res_str = ""
        
        while True:
            chunk = await c_reader.read(1024)
            if not chunk: 
                break
            
            res_str = res_str + bytes.decode(chunk)
            
            if bytes.decode(chunk).find("\r\n\r\n") > 0: 
                break

        print("<-- {0}".format(res_str))
        if res_str.lower().find("200 connection established") < 0:
            raise Exception("Unable to establish HTTP-Tunnel: {0}".format(res_str))

def main(local_ip, local_port, remote_ip, remote_port, target_ip=None, target_port=None):
    tp = tcpproxy(remote_ip, remote_port, target_ip, target_port)
    
    loop = asyncio.get_event_loop()
    coro = asyncio.start_server(tp.accept_handle, local_ip, local_port, loop=loop)
    server = loop.run_until_complete(coro)

    print('Serving on {}'.format(server.sockets[0].getsockname()))
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

if __name__ == "__main__":
#    main('0.0.0.0', 443, '127.0.0.1', 7474)